#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2022/2/10 17:26
# @Author  : CHENWang
# @Site    :
# @File    : df_tool.py
# @Software: PyCharm

"""
脚本说明:
"""

import numpy as np
import pandas as pd
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.hammer import slim


def standardize_each_col(df, **kwargs):
    """
    对 DF 的各列进行标准化
        可指定：哪些列不需要进行标准化
        可指定：标准化之后的数据覆盖源数据还是新增一列存储
    :param pd.DataFrame df: 数据
    :param kwargs:
        - except_these_cols，list of str 或 str，这些列不进行标准化，默认：[]
        - replace_or_add，str，这些列进行标准化之后，是替代原列还是新增，默认：replace
        - delete_old_cols，bool，这些列进行标准化之后，是否删除原列，默认：False
        - s_method，str，标准化方法，默认：normal
    :return: pd.DataFrame
    """
    except_these_cols = kwargs.pop('except_these_cols', [])
    replace_or_add = kwargs.pop('replace_or_add', 'replace')
    delete_old_cols = kwargs.pop('delete_old_cols', False)
    s_method = kwargs.pop('s_method', 'normal')

    slim(df, name='标准化之前的数据', a_f=1)
    cols_needed = list(df.columns)
    if isinstance(except_these_cols, str):
        except_these_cols = [except_these_cols]
    for a_col in except_these_cols:
        cols_needed.remove(a_col)

    minimum_std = 0.000000001
    for a_col in cols_needed:
        LOG.info(f'现在对列（{a_col}）进行标准化：{s_method}')
        new_col = f'z_{a_col}' if replace_or_add == 'add' else a_col  # 在新列前增加z_
        if s_method == 'normal':
            col_mean, col_std = df[a_col].mean(), df[a_col].std()
            x_4, y_4 = '均值', '标准差'
        elif s_method == 'uniform':
            col_mean, col_std = df[a_col].min(), df[a_col].max()
            x_4, y_4 = '最小值', '最大值'
        else:
            raise NotImplementedError(s_method)
        LOG.info(f'列（{a_col}）{x_4}：{col_mean}，{y_4}：{col_std}')
        if col_std <= minimum_std:
            LOG.error(f'发生列（{a_col}）{y_4}<=既定值（{minimum_std}）！！！')
        df.loc[:, new_col] = df[a_col].apply(
            lambda x: None if col_std <= minimum_std else (x - col_mean) / col_std
        )
        if delete_old_cols:
            del df[a_col]
    slim(df, name=f'标准化({s_method})之后的数据', a_f=1)
    return df


def df_to_np_array(df):
    """
    不想写太多的 import numpy as np 才写了这个函数，因为大部分时候本文件都已经被 import
    """
    my_array = np.array(df)
    return my_array


def split_a_df_into_batches(df, how_many):
    """
    把一个大的 pd.DataFrame 分成几个小的，并放入一个 list

    :raise ValueError: 如果传进来的份数=0
    :param pd.DataFrame df: 数据
    :param int how_many: 分成几份
    :return: list
    """
    if how_many == 0:
        raise ValueError('需要分成的批数不能是 0')
    total_count = df.shape[0]
    actual_batches_count = min(total_count, how_many)
    elem_count_4_a_batch = int(total_count / actual_batches_count)

    start_idx = 0
    batch_list = []
    while start_idx < total_count:
        end_idx = min(start_idx + elem_count_4_a_batch, total_count)
        batch_list.append(df.iloc[range(start_idx, end_idx)])
        start_idx = end_idx
    assert sum([x.shape[0] for x in batch_list]) == total_count
    return batch_list


def merge(df_list, key, **kwargs):
    """
    合并几个 DF 成一个，通过主键（key 表示）
        - 为什么有？如果一次性 merge 多个 df，就得写一连串的 merge，不好看。
        - 另一些代码复用。2020年3月22日16:12:29
        - drop空 的判断在 fill空 之前
    :raise ValueError: 如果传入的 df_list 不是 list
    :param df_list: 每一个 element 都是 DF
    :param str_OR_list key: 主键
    :param kwargs:
        - how，str，俩 DF join 的方式，默认：outer
        - sepper，str，如果遇到使用特殊符号填充空值，那么使用这个，默认：-
        - drop_na，bool，是不是把空值的地方drop，默认：False
        - fill_na，bool，是不是用约定好的符号填充空值的地方，默认：True
        - fill_method，str，填充数据的方法，默认：None
        - res_name，str，最终数据的名字，显示结果时加上去，默认：None
        - verbose，bool，是不是把每一步join之后的结果显示出来，默认：False
    :return: df
    """
    if not isinstance(df_list, list):
        raise ValueError('expect a list')
    how = kwargs.pop('how', 'outer')
    sepper = kwargs.pop('sepper', '-')
    drop_na = kwargs.pop('drop_na', False)
    fill_na = kwargs.pop('fill_na', True)
    fill_method = kwargs.pop('fill_method', None)
    res_name = kwargs.pop('res_name', None)
    verbose = kwargs.pop('verbose', False)

    df_num = len(df_list)
    ans = None
    for idx in range(df_num):
        a_df = df_list[idx]
        if not isinstance(a_df, pd.DataFrame):
            raise RuntimeError('the %sth element is not a pd.DataFrame: %s' %
                               (idx + 1, a_df))
        # 這個方法有點扯，如果有一列原本是int，merge完之後，會變成float
        ans = a_df if ans is None else pd.merge(
            ans, a_df, right_on=key, left_on=key, how=how
        )
        if verbose:
            slim(ans, name='经过第%s次%s join后' % (idx + 1, how), a_f=True)

            # --DEBUG--START--
            # if idx + 1 == 2:
            #     slim(ans, a_f=0, show_index=1, show_all=1,
            #          name='经过第%s次%s join后' % (idx+1, how))
            # --DEBUG--END--

    if drop_na:
        ans.dropna(inplace=True)
    if fill_na:
        if fill_method == 'ffill':
            ans.fillna(method=fill_method, inplace=True)
        else:
            ans.fillna(sepper, inplace=True)
    msg = f'{df_num}个DF {how} merge 结果，填充na:{fill_na}'
    msg += '' if not fill_na else f'，填充符:{sepper}，填充方法:{fill_method}'
    if res_name:
        msg = '$' + res_name + '$ ' + msg
    slim(ans, a_f=1, name=msg)
    return ans


def single_row_df_to_dict(row):
    """
    把某一个只有一行的 df 转成 dict
        如果该 df 为空，则结果可能很怪
    :param pd.DataFrame row: 那个只有一行的 df
    :return: dict
    """
    keys = row.columns.values
    values = row.values[0]
    ans = {}
    for ii in range(len(keys)):
        ans.update({keys[ii]: values[ii]})
    return ans


def cpt_group_sort(data, group_by, sort_by, how_many, **kwargs):
    """
    针对 df 的组内排序

    :param pd.DataFrame data: 数据
    :param str group_by: 此列分组
    :param str sort_by: 此列排序
    :param how_many: 组内前几
    :param kwargs:
        - if_debugging，bool，是否处于调试中（多输出一些中间结果），默认：False
    :return: pd.DataFrame
    """
    if_debugging = kwargs.pop('if_debugging', False)

    df = data.copy()  # 2020年2月20日16:40:09，免得无意之中修改了本想再次使用的数据
    if df.shape[0] != 0:
        df['group_sort'] = df[sort_by] \
            .groupby(by=df[group_by]) \
            .rank(ascending=0, method='dense')  # 逆序，密集的
        if if_debugging:
            slim(df, name='经过排序之后', a_f=1)
        df = df[df['group_sort'] < (how_many + 1)]  # 这里要 +1
        if if_debugging:
            slim(df, name='经过筛选之后', a_f=1)
        df.index = range(df.shape[0])
        df = df[[sort_by, group_by]]
        df = df.groupby(by=group_by, as_index=False).sum()
    else:
        df = df[[sort_by, group_by]]
    df.rename(columns={sort_by: f'sum_of_first_{how_many}'}, inplace=True)
    if if_debugging:
        slim(df, a_f=1, show_index=1,
             name=f'按 {group_by} 分组，按 {sort_by} 排序后，前 {how_many} 之和')

    return df


def calendar_reindex(df, end=None, fmt='%Y-%m-%d'):
    """
    以自然日作为index
    :param df:
    :return: df
    """
    range_end = end if end else df.index.max()
    calendar = pd.date_range(df.index.min(), range_end).strftime(fmt)
    out_df = df.reindex(calendar)
    return out_df


def reduce_mem_usage(df):
    """
    iterate through all the columns of a dataframe and modify the data type to reduce memory usage.
    """

    # This function is used to reduce memory of a pandas dataframe
    # The idea is cast the numeric type to another more memory-effective type
    # For ex: Features "age" should only need type='np.int8'
    # Source: https://www.kaggle.com/gemartin/load-data-reduce-memory-usage
    # @from: https://www.kaggle.com/arjanso/reducing-dataframe-memory-size-by-65/code

    # 计算当前内存
    start_mem_usg = df.memory_usage().sum() / 1024 ** 2
    print("Memory usage of the dataframe is :", start_mem_usg, "MB")

    for col in df.columns:
        col_type = df[col].dtype
        # 这里只过滤了objectd格式，如果你的代码中还包含其他类型，请一并过滤
        if col_type != object and col_type.name != 'category' and 'datetime' not in col_type.name:
            print("**************************")
            print("columns: ", col)
            print("dtype before", col_type.name)

            mmax = df[col].max()
            mmin = df[col].min()
            # # Integer does not support NA, therefore Na needs to be filled
            # if not np.isfinite(df[col]).all():
            #     NAlist.append(col)
            #     df[col].fillna(-999, inplace=True)  # 用-999填充

            if str(col_type)[:3] == 'int':
                # make interger / unsigned Integer datatypes
                if mmin >= 0:  # 最小值大于0，转换成无符号整型
                    if mmax <= np.iinfo(np.uint8).max:
                        df[col] = df[col].astype(np.uint8)
                    elif mmax <= np.iinfo(np.uint16).max:
                        df[col] = df[col].astype(np.uint16)
                    elif mmax <= np.iinfo(np.uint32).max:
                        df[col] = df[col].astype(np.uint32)
                    else:
                        df[col] = df[col].astype(np.uint64)
                else:  # 转换成有符号整型
                    if mmin > np.iinfo(np.int8).min and mmax < np.iinfo(np.int8).max:
                        df[col] = df[col].astype(np.int8)
                    elif mmin > np.iinfo(np.int16).min and mmax < np.iinfo(np.int16).max:
                        df[col] = df[col].astype(np.int16)
                    elif mmin > np.iinfo(np.int32).min and mmax < np.iinfo(np.int32).max:
                        df[col] = df[col].astype(np.int32)
                    elif mmin > np.iinfo(np.int64).min and mmax < np.iinfo(np.int64).max:
                        df[col] = df[col].astype(np.int64)
            else:
                if mmin > np.finfo(np.float16).min and mmax < np.finfo(np.float16).max:
                    df[col] = df[col].astype(np.float16)
                elif mmin > np.finfo(np.float32).min and mmax < np.finfo(np.float32).max:
                    df[col] = df[col].astype(np.float32)
                else:
                    df[col] = df[col].astype(np.float64)

        elif 'datetime' not in col_type.name:
            df[col] = df[col].astype('category')

        print("dtype after", df[col].dtype)
        print("********************************")

    print("___MEMORY USAGE AFTER COMPLETION:___")
    mem_usg = df.memory_usage().sum() / 1024 ** 2
    print("Memory usage is: ", mem_usg, " MB")
    print("This is ", 100 * mem_usg / start_mem_usg, "% of the initial size")

    return df


### dataframe 根据列值筛选 ###
def df_filter_test():
    df = pd.DataFrame(np.random.uniform(0, 10, (100, 3)), columns=['A', 'B', 'C'])
    # 如果我们只需要‘B’列小于5的行：当然>,<,==,>=,<=都是相同的道理。小心“等于”一定是用‘==’，如果用‘=’就不是判断大小了。
    res = df[df['B'] < 5]
    # 如果有多个条件，就用&将多个条件连接起来，每个条件用（）括起来
    res = df[(df['B'] < 5) & (df['C'] > 5)]
    # ”&“是要求两个条件都要满足，是交集。也可以“|”，只要满足一个条件都会被筛选出来，是并集。
    res = df[(df['B'] < 5) | (df['C'] > 5)]
    # 通过B筛选，获取A.C列数据
    res = df[['A', 'C']][df['B'] < 5]

    # 除了比较大小，还可以如果要选择某列等于多个数值或者字符串时，要用到.isin(), isin()括号里面应该是个list
    res = df[df['B'].isin([1, 5])]
    # 其实平时用的最最多的筛选，应该是字符串的模糊筛选，在SQL语音里用的是like。在pandas里面我们可以用.str.contains()
    str_df = df.astype(str)
    res = str_df[str_df['B'].str.contains('123')]
    # 当然也可以用‘|’进行多个条件筛选：注意，这个‘|’是在引号内的，而不是将两个字符串分别引起来。’&‘在这里不能用。
    res = str_df[str_df['B'].str.contains('122|123')]
    res = str_df[str_df['B'].str.contains('122') & str_df['B'].str.contains('146')]
    # 想获取不等于的数据
    int_df = df.astype(int)
    res = int_df[~int_df['B'] == 5]


def df_duplicate_test():
    df = pd.DataFrame(np.random.uniform(0, 10, (100, 3)), columns=['A', 'B', 'C'])
    int_df = df.astype(int)
    df.set_index('B', inplace=True)
    res = df[~df.index.duplicated(keep='first')]  # index去重
    res = df.drop_duplicates(subset=['A', 'C'], inplace=True)  # 根据列值去除重复行


if __name__ == '__main__':
    df_filter_test()
